from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy import event, text
from typing import AsyncGenerator
import asyncio
import time
import re

try:
    import structlog

    logger = structlog.get_logger(__name__)
except ImportError:
    import logging

    logger = logging.getLogger(__name__)

from .config import settings


class Base(DeclarativeBase):
    """
    SQLAlchemy声明式基类
    所有数据库模型都应该继承此类
    """

    pass


class DatabaseManager:
    """
    数据库连接管理器
    负责管理数据库连接池和会话
    """

    def __init__(self):
        self.engine = None
        self.async_session_maker = None
        self._initialized = False
        self._heartbeat_task = None

    async def initialize(self) -> None:
        """
        初始化数据库连接
        创建异步引擎和会话工厂
        """
        if self._initialized:
            logger.warning("数据库已经初始化，跳过重复初始化")
            return

        try:
            # 创建异步数据库引擎
            self.engine = create_async_engine(
                settings.database_url,
                echo=False,  # 关闭SQLAlchemy默认日志，使用我们的美观日志
                pool_size=10,  # 连接池大小
                max_overflow=20,  # 最大溢出连接数
                pool_pre_ping=True,  # 连接前检查连接有效性
                pool_recycle=settings.db_heartbeat_interval
                * 2,  # 连接回收时间应大于心跳间隔
            )

            # 添加SQL查询拦截器
            self._setup_sql_interceptor()

            # 创建异步会话工厂
            self.async_session_maker = async_sessionmaker(
                bind=self.engine,
                class_=AsyncSession,
                expire_on_commit=False,
                autoflush=True,
                autocommit=False,
            )

            # 测试数据库连接
            await self._test_connection()

            self._initialized = True

            logger.info(
                "数据库连接初始化成功",
                database_url=settings.database_url.replace(
                    settings.postgres_password, "***"
                ),
                pool_size=10,
                max_overflow=20,
            )

        except Exception as e:
            logger.error(
                "数据库连接初始化失败",
                error=str(e),
                database_url=settings.database_url.replace(
                    settings.postgres_password, "***"
                ),
            )
            raise

    def _setup_sql_interceptor(self) -> None:
        """
        设置SQL查询拦截器，用于记录美观的SQL执行日志
        """

        @event.listens_for(self.engine.sync_engine, "before_cursor_execute")
        def receive_before_cursor_execute(
            conn, cursor, statement, parameters, context, executemany
        ):
            context._query_start_time = time.time()

            # 格式化SQL语句
            formatted_sql = self._format_sql(statement)

            logger.info(
                "🔍 执行SQL查询",
                sql=formatted_sql,
                parameters=self._safe_format_parameters(parameters),
                executemany=executemany,
            )

        @event.listens_for(self.engine.sync_engine, "after_cursor_execute")
        def receive_after_cursor_execute(
            conn, cursor, statement, parameters, context, executemany
        ):
            total_time = time.time() - context._query_start_time

            logger.info(
                "✅ SQL查询完成",
                execution_time_ms=round(total_time * 1000, 2),
                rowcount=cursor.rowcount if hasattr(cursor, "rowcount") else "N/A",
            )

    def _format_sql(self, sql: str) -> str:
        """
        格式化SQL语句，使其更易读
        """
        # 移除多余的空白字符
        sql = re.sub(r"\s+", " ", sql.strip())

        # 截断过长的SQL语句
        if len(sql) > 200:
            sql = sql[:200] + "..."

        return sql

    def _safe_format_parameters(self, parameters) -> str:
        """
        安全地格式化SQL参数，避免敏感信息泄露
        """
        if not parameters:
            return "None"

        try:
            # 如果参数是字典或元组，转换为字符串并截断
            param_str = str(parameters)
            if len(param_str) > 100:
                param_str = param_str[:100] + "..."
            return param_str
        except Exception:
            return "<无法格式化参数>"

    async def _test_connection(self) -> None:
        """
        测试数据库连接是否正常
        """
        try:
            async with self.engine.begin() as conn:
                result = await conn.execute(text("SELECT 1 as test_connection"))
                test_result = result.scalar()

                if test_result == 1:
                    logger.info("数据库连接测试成功")
                else:
                    raise Exception("数据库连接测试失败：返回值不正确")

        except Exception as e:
            logger.error("数据库连接测试失败", error=str(e))
            raise

    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        """
        获取数据库会话
        使用异步上下文管理器确保会话正确关闭
        """
        if not self._initialized:
            await self.initialize()

        async with self.async_session_maker() as session:
            try:
                yield session
            except Exception as e:
                await session.rollback()
                logger.error("数据库会话异常，已回滚", error=str(e))
                raise
            finally:
                await session.close()

    async def close(self) -> None:
        """
        关闭数据库连接
        """
        if self.engine:
            # 停止心跳任务
            await self._stop_heartbeat()

            await self.engine.dispose()
            logger.info("数据库连接已关闭")
            self._initialized = False

    async def _reconnect_database(self) -> bool:
        """
        重新连接数据库
        Returns:
            bool: 重连是否成功
        """
        if not settings.db_reconnect_enabled:
            logger.info("数据库重连功能已禁用")
            return False

        logger.info("尝试重新连接数据库")

        try:
            # 关闭现有连接
            if self.engine:
                await self.engine.dispose()

            # 重新初始化数据库连接
            await self.initialize()
            logger.info("数据库重连成功")
            return True

        except Exception as e:
            logger.error("数据库重连失败", error=str(e), error_type=type(e).__name__)
            return False

    async def _run_heartbeat(self) -> None:
        """
        运行数据库心跳检测任务
        定期执行简单查询以保持连接活跃
        """
        if not settings.db_heartbeat_enabled:
            logger.info("数据库心跳功能已禁用")
            return

        logger.info(
            "启动数据库心跳检测任务", interval=f"{settings.db_heartbeat_interval}秒"
        )

        while True:
            try:
                # 等待心跳间隔
                await asyncio.sleep(settings.db_heartbeat_interval)

                # 检查数据库引擎是否已初始化
                if not self.engine or not self._initialized:
                    logger.debug("数据库未初始化，跳过心跳检测")
                    continue

                # 执行心跳查询
                async with self.engine.begin() as conn:
                    result = await conn.execute(text("SELECT 1"))
                    heartbeat_result = result.scalar()

                    if heartbeat_result == 1:
                        logger.debug("数据库心跳检测成功")
                    else:
                        logger.warning(
                            "数据库心跳检测返回异常值", result=heartbeat_result
                        )

            except asyncio.CancelledError:
                logger.info("数据库心跳任务已取消")
                break
            except Exception as e:
                logger.error(
                    "数据库心跳检测失败", error=str(e), error_type=type(e).__name__
                )

                # 如果启用了重连功能，尝试重新连接数据库
                if settings.db_reconnect_enabled:
                    reconnect_success = False
                    for attempt in range(settings.db_reconnect_max_retries):
                        try:
                            logger.info(f"尝试第{attempt + 1}次重连数据库")
                            reconnect_success = await self._reconnect_database()
                            if reconnect_success:
                                logger.info(f"数据库重连成功，重试次数: {attempt + 1}")
                                break
                            else:
                                # 等待重试间隔
                                if attempt < settings.db_reconnect_max_retries - 1:
                                    logger.info(
                                        f"重连失败，{settings.db_reconnect_retry_interval}秒后进行下次重试"
                                    )
                                    await asyncio.sleep(
                                        settings.db_reconnect_retry_interval
                                    )
                        except Exception as reconnect_error:
                            logger.error(
                                f"第{attempt + 1}次重连尝试失败",
                                error=str(reconnect_error),
                                error_type=type(reconnect_error).__name__,
                            )
                            if attempt < settings.db_reconnect_max_retries - 1:
                                await asyncio.sleep(
                                    settings.db_reconnect_retry_interval
                                )

                    if not reconnect_success:
                        logger.error(
                            "数据库重连失败，已达到最大重试次数",
                            max_retries=settings.db_reconnect_max_retries,
                        )
                else:
                    logger.info("数据库重连功能已禁用，跳过重连尝试")

    async def _start_heartbeat(self) -> None:
        """
        启动数据库心跳任务
        """
        if not settings.db_heartbeat_enabled:
            return

        # 取消已存在的任务
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
            except asyncio.CancelledError:
                pass

        # 启动新任务
        self._heartbeat_task = asyncio.create_task(self._run_heartbeat())
        logger.info("数据库心跳任务已启动")

    async def _stop_heartbeat(self) -> None:
        """
        停止数据库心跳任务
        """
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
                logger.info("数据库心跳任务已停止")
            except asyncio.CancelledError:
                logger.info("数据库心跳任务已取消")


# 全局数据库管理器实例
db_manager = DatabaseManager()


# 依赖注入函数，用于FastAPI路由
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """
    FastAPI依赖注入函数
    用于在路由中获取数据库会话

    使用方法:
    @app.get("/example")
    async def example_route(db: AsyncSession = Depends(get_db)):
        # 使用db进行数据库操作
        pass
    """
    async for session in db_manager.get_session():
        yield session
